import unittest
import torch
import torchvision.transforms as transforms
from matplotlib import pyplot as plt

import mymodel


def test_loadmodel():
    path_model = "./saveModel/lenet5-model.pkl"
    net_load = torch.load(path_model)
    print(net_load)
    from PIL import Image

    I = Image.open('./saveModel/8.png')
    I = I.resize((28, 28), Image.ANTIALIAS)
    # convert 支持 1，L，P，RGB，RGBA，CMYK，YCbCr，I，F， 其中L是转成灰度
    L = I.convert('L')
    plt.imshow(L, cmap='gray')
    plt.show()

    # 操作打包
    transform = transforms.Compose([
        transforms.ToTensor()
        , transforms.Normalize((0.1037,), (0.3081,))
    ])

    print(L)

    im = transform(L)  # [C, H, W]
    plt.imshow(im)
    plt.show()
    print('-' * 50)
    print(im)
    im = torch.unsqueeze(im, dim=0)  # [N, C, H, W]
    print('-' * 50)
    print(im)

    with torch.no_grad():
        outputs = net_load(im)
        _, predict = torch.max(outputs.data, 1)
        print(predict)

def test_loadmodelstate():
    path_model = "./saveModel/lenet5-model_state_dict.pkl"
    netstate_load = torch.load(path_model)
    net_load = mymodel.LeNet5()
    net_load.load_state_dict(state_dict = netstate_load)
    from PIL import Image

    I = Image.open('./saveModel/8.png')
    I = I.resize((28, 28), Image.ANTIALIAS)
    # convert 支持 1，L，P，RGB，RGBA，CMYK，YCbCr，I，F， 其中L是转成灰度
    L = I.convert('L')
    plt.imshow(L, cmap='gray')
    plt.show()

    # 操作打包
    transform = transforms.Compose([
        transforms.ToTensor()
        , transforms.Normalize((0.1037,), (0.3081,))
    ])

    print(L)

    im = transform(L)  # [C, H, W]
    im = torch.unsqueeze(im, dim=0)  # [N, C, H, W]


    with torch.no_grad():
        outputs = net_load(im)
        _, predict = torch.max(outputs.data, 1)
        print(predict)


if __name__ == '__main__':
    test_loadmodelstate()
